# -*- coding: utf-8 -*-
# Author   : ZhangQing
# Time     : 2025-07-15 23:11
# File     : polygon_adapter.py
# Project  : dynamic-portfolio-optimizer
# Desc     :
# src/data/adapters/polygon_adapter.py
from polygon import RESTClient
from typing import Dict, List, Any
import pandas as pd
from datetime import datetime, timedelta
from .base_adapter import BaseDataAdapter, RateLimitDecorator


class PolygonAdapter(BaseDataAdapter):
    """Polygon.io数据适配器"""

    def __init__(self, config):
        super().__init__(config)
        self.client = RESTClient(config.api_key)

    @RateLimitDecorator(calls=5, period=60)
    def get_stock_data(self, symbol: str, start_date: datetime,
                       end_date: datetime, interval: str = '1d') -> pd.DataFrame:
        """获取股票数据"""
        try:
            # 转换时间间隔
            timespan_map = {
                '1m': 'minute',
                '5m': 'minute',
                '1h': 'hour',
                '1d': 'day'
            }

            multiplier_map = {
                '1m': 1,
                '5m': 5,
                '1h': 1,
                '1d': 1
            }

            timespan = timespan_map.get(interval, 'day')
            multiplier = multiplier_map.get(interval, 1)

            # 获取数据
            bars = self.client.get_aggs(
                ticker=symbol,
                multiplier=multiplier,
                timespan=timespan,
                from_=start_date.strftime('%Y-%m-%d'),
                to=end_date.strftime('%Y-%m-%d')
            )

            if not bars:
                self.logger.warning(f"Polygon: 未获取到{symbol}的数据")
                return pd.DataFrame()

            # 转换为DataFrame
            data_list = []
            for bar in bars:
                data_list.append({
                    'timestamp': datetime.fromtimestamp(bar.timestamp / 1000),
                    'open': bar.open,
                    'high': bar.high,
                    'low': bar.low,
                    'close': bar.close,
                    'volume': bar.volume
                })

            data = pd.DataFrame(data_list)
            data.set_index('timestamp', inplace=True)
            data['source'] = 'polygon'
            data['symbol'] = symbol

            return self.standardize_data(data)

        except Exception as e:
            self.logger.error(f"Polygon获取股票数据失败: {e}")
            return pd.DataFrame()

    @RateLimitDecorator(calls=5, period=60)
    def get_options_data(self, symbol: str, expiration_date: datetime) -> pd.DataFrame:
        """获取期权数据"""
        try:
            # 获取期权链
            options_contracts = self.client.list_options_contracts(
                underlying_ticker=symbol,
                expiration_date=expiration_date.strftime('%Y-%m-%d')
            )

            if not options_contracts:
                self.logger.warning(f"Polygon: 未获取到{symbol}的期权数据")
                return pd.DataFrame()

            options_data = []
            for contract in options_contracts:
                # 获取期权详情
                try:
                    details = self.client.get_options_contract(contract.ticker)

                    options_data.append({
                        'symbol': symbol,
                        'contract_symbol': contract.ticker,
                        'strike': details.strike_price,
                        'expiration': expiration_date,
                        'type': 'call' if details.contract_type == 'call' else 'put',
                        'source': 'polygon'
                    })
                except:
                    continue

            return pd.DataFrame(options_data)

        except Exception as e:
            self.logger.error(f"Polygon获取期权数据失败: {e}")
            return pd.DataFrame()

    @RateLimitDecorator(calls=5, period=60)
    def get_fundamentals(self, symbol: str) -> Dict[str, Any]:
        """获取基本面数据"""
        try:
            # 获取公司详情
            ticker_details = self.client.get_ticker_details(symbol)

            if not ticker_details:
                self.logger.warning(f"Polygon: 未获取到{symbol}的基本面数据")
                return {}

            fundamentals = {
                'symbol': symbol,
                'company_name': ticker_details.name,
                'sector': getattr(ticker_details, 'sector', ''),
                'industry': getattr(ticker_details, 'industry', ''),
                'market_cap': getattr(ticker_details, 'market_cap', 0),
                'shares_outstanding': getattr(ticker_details, 'share_class_shares_outstanding', 0),
                'description': getattr(ticker_details, 'description', ''),
                'source': 'polygon'
            }

            return fundamentals

        except Exception as e:
            self.logger.error(f"Polygon获取基本面数据失败: {e}")
            return {}

    @RateLimitDecorator(calls=5, period=60)
    def get_news(self, symbol: str, limit: int = 10) -> List[Dict[str, Any]]:
        """获取新闻数据"""
        try:
            # 获取新闻
            news_items = self.client.list_ticker_news(
                ticker=symbol,
                limit=limit
            )

            processed_news = []
            for item in news_items:
                processed_news.append({
                    'symbol': symbol,
                    'title': item.title,
                    'summary': item.description,
                    'url': item.article_url,
                    'published_at': datetime.fromisoformat(item.published_utc.replace('Z', '+00:00')),
                    'source': 'polygon',
                    'publisher': item.publisher.name if item.publisher else ''
                })

            return processed_news

        except Exception as e:
            self.logger.error(f"Polygon获取新闻数据失败: {e}")
            return []
